from typing import Dict, List

from .node_handler import NodeHandler
from .registry import operator_registry
from ..generator import TensorConstructorGenerator,StrategyGenerator
from ..shard.placement_types import TensorMeta,OperationData,OperationDataType
__all__ = ['TensorConstructorHandler']


# @operator_registry.register('arange')
class TensorConstructorHandler(NodeHandler):
    """
    A TensorConstructorHandler which deals with the sharding strategies for tensor constructor operations, such as torch.arange.
    """

    def get_strategy_generator(self) -> List[StrategyGenerator]:
        op_data_mapping = self.get_operation_data_mapping()
        generators = []
        generators.append(TensorConstructorGenerator(op_data_mapping, self.device_mesh,self.node.name))
        return generators

    def get_operation_data_mapping(self) -> Dict[str, OperationData]:
        output_data = TensorMeta(self.node.meta['tensor_meta'].shape,self.node.meta['tensor_meta'].stride,self.node.meta['tensor_meta'].dtype)
        physical_output_operand = OperationData(name=str(self.node), data=output_data,type=OperationDataType.OUTPUT)

        mapping = {"output": physical_output_operand}

        return mapping
